from matplotlib import pyplot as plt
from matplotlib.figure import Figure
from matplotlib.axes import Axes
from cartopy import crs as ccrs
import numpy as np
import xarray as xr
from matplotlib.contour import QuadContourSet
from matplotlib.gridspec import GridSpec
from . import cal

def basemap(fig: Figure, mapRange: list, nrows: int = 1, ncols: int = 1,
            index: int = 1, proj: ccrs.Projection = None,
            central_longitude: float = 0,
            coastlines: bool | dict = {"color": "black", "linewidth": 0.5},
            gridlines: bool | dict = {"color": "k", "linestyle": (0, (5, 5)),
                                      "linewidth": 0.5},
            lon_ticks: list | None = None, lat_ticks: list | None = None,
            ):
    """
    以 subplot 画基本地图底图

    ----------------------------------

    @param fig: Figure
    @param mapRange: 地图范围，[lon_min, lon_max, lat_min, lat_max]
    @param nrows: int, optional, add_subplot 的参数
    @param ncols: int, optional, add_subplot 的参数
    @param index: int, optional, add_subplot 的参数
    @param proj: ccrs.Projection, optional, 投影方式，默认为 PlateCarree
    @param central_longitude: float, optional, proj=None 时 PlateCarree 投影方式
        的中心经度，默认为 0。
    @param coastlines: bool | dict, optional, 是否画海岸线，如果为 dict，则作为
        ax.coastlines 的参数传入
    @param lon_ticks: list | None, optional, 指定经度刻度，不设置则5度间隔
    @param lat_ticks: list | None, optional, 指定纬度刻度，不设置则5度间隔

    """
    if proj is None:
        proj = ccrs.PlateCarree(central_longitude=central_longitude)

    ax = fig.add_subplot(nrows, ncols, index, projection=proj)
    ax.set_extent(mapRange, crs=proj)

    if coastlines:
        ax.coastlines(**coastlines)

    if lon_ticks is None:
        lon_ticks = range(int(mapRange[0]), int(mapRange[1])+1, 5)
    if lat_ticks is None:
        lat_ticks = range(int(mapRange[2]), int(mapRange[3])+1, 5)

    if gridlines:
        gl = ax.gridlines(draw_labels=True, xlocs=lon_ticks, ylocs=lat_ticks,
                          **gridlines)
        gl.top_labels = False
        gl.right_labels = False

    return ax

def quiver(x: xr.DataArray | np.ndarray, y: xr.DataArray | np.ndarray,
           u: xr.DataArray | np.ndarray, v: xr.DataArray | np.ndarray, 
           ax=None, proj: ccrs.Projection = ccrs.PlateCarree(),
           scale: float = 1, units: str = 'xy', width: float = 0.1,
           lowRes: int = 4, zorder:int = 5, **kwargs):
    
    if ax is None:
        ax = plt.gca()
    
    q = ax.quiver(x[::lowRes], y[::lowRes], 
                  u[::lowRes, ::lowRes], v[::lowRes, ::lowRes], 
                  transform=proj, scale=scale, units=units,
                  width=width, zorder=zorder, **kwargs)
    return q


def streamplot(x: xr.DataArray | np.ndarray, y: xr.DataArray | np.ndarray,
           u: xr.DataArray | np.ndarray, v: xr.DataArray | np.ndarray, 
           ax=None, proj: ccrs.Projection = ccrs.PlateCarree(),
           density=1, linewidth=0.7, color='k', zorder:int = 5, **kwargs):
    
    if ax is None:
        ax = plt.gca()
    
    q = ax.streamplot(x, y, u, v, 
                  transform=proj,density=density, linewidth=linewidth,
                    zorder=zorder, color=color, **kwargs)
    return q

def quiverkey(q, x: float = 0.92, y: float = 0.97, length: float = 1,
               unit:str = 'm/s', fontsize: int = 12):
    plt.quiverkey(q, x, y, length, "%s %s" % (length, unit), labelpos='E', 
                  coordinates='figure', fontproperties={'size': fontsize})

def colorbar(fig: Figure, axesRange: list, C, horizontal: bool = True, **kwargs):

    ax = fig.add_axes(axesRange)
    cb = plt.colorbar(C, cax=ax, 
                 orientation='horizontal' if horizontal else 'vertical',
                 **kwargs)
    return cb

def subplot_label(text: str, x: float = 0.02, y: float = 0.97, 
                  ax=None, bg: bool = True, bgPad=1, zorder=100, fontsize=12,
                  ha='left', va='top'):
    if ax is None:
        ax = plt.gca()

    if bg:
        ax.text(x, y, text, transform=ax.transAxes, bbox=dict(
            facecolor='white', edgecolor='none', pad=bgPad
            ), ha=ha, va=va, zorder=zorder, fontsize=fontsize)
    else:
        ax.text(x, y, text, transform=ax.transAxes, ha=ha, va=va, 
                zorder=zorder, fontsize=fontsize)

def part_of_cmap(cmap, start: float = 0, end: float = 1):
    """
    裁剪自带的 colormap 的一部分，生成新的 colormap

    """
    
    from matplotlib import cm
    from matplotlib.colors import LinearSegmentedColormap
    
    cmap = cm.get_cmap(cmap)
    colors = cmap(np.linspace(start, end, 256))
    new_cmap = LinearSegmentedColormap.from_list('part_of_%s' % cmap.name, colors)
    return new_cmap


def fill_land(color: str = '#888888',zorder: int = 0):
    """
    陆地填色
    """

    from cartopy.feature import LAND
    ax = plt.gca()
    ax.add_feature(LAND, facecolor=color, zorder=zorder)


def subplot_gl(ax, top=False, right=False, left=False, bottom=False,
               linestyle='--', linewidth=0.5, color='gray',
               xlocs=[-160, 120, 140, 160, 180],
               ylocs=[-20, -10, 0, 10, 20], **kwargs):
    """
    邮票图画经纬度网格
    根据情况设置是否画上下左右的经纬度标签
    """
    gl = ax.gridlines(crs=ccrs.PlateCarree(0), draw_labels=True,
                    linewidth=linewidth, color=color, linestyle=linestyle,
                    xlocs=xlocs, ylocs=ylocs, **kwargs)
    gl.top_labels = top
    gl.right_labels = right
    gl.left_labels = left
    gl.bottom_labels = bottom

def set_hatch_color(contourf: QuadContourSet, color: str | list[str],
                     borderWidth=0):
    """
    设置用 contourf 画的 hatch 的颜色

    @param contourf: contourf 的返回对象。
    @param color: hatch 的颜色，可以是单个颜色，也可以是一个 list，
        对应每个 hatch 的颜色。
    @param borderWidth: hatch 外边界宽度，一般设置为0隐藏。
    """

    for i, collection in enumerate(contourf.collections):
        if isinstance(color, list):
            collection.set_edgecolor(color[i])
        else:
            collection.set_edgecolor(color)

    for collection in contourf.collections:
        collection.set_linewidth(borderWidth)


def grid_plot(row: int, col: int, ratio: float = 1.5, 
              figsize: tuple = None, figZoom: float = 1,
              sub_adj: tuple = (0.05, 0.95, 0.05, 0.95),
              wSpace: float = 0.05, hSpace: float = 0.05,
              returnGS: bool = False, subplot={}, **kwargs) -> tuple[Figure, list[Axes] | GridSpec]:
    """
    画阵列多子图框架，使用时用 axes 制图：
    ```
    row = 5
    col = 5
    fig, axes = grid_plot(row, col)

    for i in range(row):
        for j in range(col):
            ax = axes[i, j]
            ax.plot(yourData)

    ```

    Parameters
    ----------
    row : int
        行数
    col : int
        列数
    ratio : float
        子图的长宽比，用于自动确定figsize，figsize=None时有效
    figsize : tuple, optional
        图片大小，None时自动计算, by default None
    figZoom : float, optional
        图片放大倍数, by default 1，每一张子图为 figZoom(inch) 高。
    sub_adj : tuple, optional
        子图间距，(left, right, bottom, top), 同 subplots_adjust, 
        by default (0.05, 0.95, 0.05, 0.95)
    wSpace : float, optional
        子图水平间距, by default 0.05
    hSpace : float, optional
        子图垂直间距, by default 0.05
    subplot : dict, optional
        子图参数, by default 
    returnGS : bool, optional
        是否返回GridSpec对象，否则返回subplot数组。 by default False
    
    Returns
    -------
    fig : matplotlib.figure.Figure
        图片对象
    axes : list | matplotlib.gridspec.GridSpec
        子图数组或GridSpec对象

    """

    if figsize is None:
        # 绘图区高度
        fig_ratio = col*ratio/row
        figHeight = row * figZoom / (sub_adj[3] - sub_adj[2])
        figWidth = row * figZoom * fig_ratio / (sub_adj[1] - sub_adj[0])
        figsize = (figWidth, figHeight)

    fig = plt.figure(figsize=figsize)
    fig.subplots_adjust(left=sub_adj[0], right=sub_adj[1],
                        bottom=sub_adj[2], top=sub_adj[3])
    gs = GridSpec(row, col, figure=fig, wspace=wSpace, hspace=hSpace, **kwargs)

    if returnGS:
        return fig, gs

    axes = []
    for i in range(row):
        axes.append([])
        for j in range(col):
            axes[i].append(fig.add_subplot(gs[i, j], **subplot))
    
    return fig, axes

def emphasis_contour(ax, x, y, data, level: list[float] | np.ndarray,
                     emphasisLevel: float | list[float] | np.ndarray, 
                     colors='k', emphasisColor='r', linewidths=0.8,
                     emphasisLinewidth=1.3, contourArgs={}, clabelArgs={}) -> None:
    """
    画等值线，并强调其中某一些

    Parameters
    ----------
    ax: axes
        当前 axes
    x: array-liked
        contour的参数
    y: array-liked
        contour的参数
    data: array-liked
        contour的参数
    level: array-liked
        要画的等值线的值，包含要强调的值
    emphasisLevel: float | array-liked
        要强调的等值线的值
    colors: matplotlib colors
        非强调的等值线的颜色
    emphasisColor: matplotlib colors
        强调的等值线的颜色
    linewidths: float
        非强调的等值线宽度
    emphasisLinewidth: float
        强调的等值线宽度
    contourAugs: dict
        contour() 函数的其他参数，以字典形式传入
    clabel: dict
        clabel() 函数的其他参数，以字典形式传入

    """
    level = np.array(level)

    if not isinstance(emphasisLevel, list) and \
        not isinstance(emphasisLevel, np.ndarray):
        emphasisLevel = [emphasisLevel]

    emphasisLevel = np.array(emphasisLevel)
    level = np.delete(level, np.where(np.isin(level, emphasisLevel)))
    
    C1 = ax.contour(x, y, data, level, colors=colors, linewidths=linewidths, **contourArgs)

    C2 = ax.contour(x, y, data, emphasisLevel, colors=emphasisColor, 
               linewidths=emphasisLinewidth, **contourArgs)
    
    for C in [C1, C2]:
        clabel = ax.clabel(C, zorder=50, **clabelArgs)
        for c in clabel:
            c.set_bbox(dict(facecolor='white', edgecolor='None', pad=0.1))
    
def subplot_tick(xtick: list[float] | np.ndarray, ytick: list[float] | np.ndarray,
                 showYTickLabel: bool = True, showXTickLabel: bool = True, ax=None,
                 xAppend: str=None, yAppend: str=None, **kwargs):
    """
    多子图画坐标轴刻度，只在边缘画。例如

    ```
    for i in range(3):
        for j in range(4):
            ax = fig.add_subplots(3, 4, i*4+j)

            # 只在子图阵列的左边和下边显示 tick 的数值
            subplot_tick([0, 1, 2], [0, 1, 2],
                         showYTickLabel = j == 0,
                         showXTickLabel = i == 2)

    ```
    Parameters
    ----------
    xtick: array-liked
        x刻度值
    ytick: array-liked
        y刻度值
    showYTickLabel: bool
        是否显示y刻度的值
    showXTickLabel: bool
        是否显示x刻度的值
    ax: axes
        axes
    xAppend: str
        xtick的后缀，例如表示经度时可以加 "E"
    yAppend: str
        ytick的后缀

    """
    if ax is None:
        ax = plt.gca()

    if xAppend is not None:
        xtickLabel = [str(i)+xAppend for i in xtick]
    
    if yAppend is not None:
        ytickLabel = [str(i)+yAppend for i in ytick]

    ax.set_xticks(xtick)
    ax.set_xticklabels(xtickLabel if showXTickLabel else [], **kwargs)
    ax.set_yticks(ytick)
    ax.set_yticklabels(ytickLabel if showYTickLabel else [], **kwargs)


def grid_colorbar(fig: Figure, axesRange: list, number: int, horizontal=False, space=0.05):
    """
    使用 grid_plot 时，出现每一行或每一列需要一个colorbar时使用。
    生成多个axes给colorbar使用。

    Parameters
    ----------
    fig: Figure

    axesRange: array-liked
        colorbar列表范围，分别为[左下角x坐标，左下角y坐标，宽度，高度]
    number: int
        colorbar数量
    horizontal: bool, default: False
        colorbar横向还是竖向分布
    space: float
        间隔

    Return
    ------
    axes: list[axes]

    """

    from matplotlib.gridspec import GridSpec

    if horizontal:
        gs = GridSpec(1, number, fig, left=axesRange[0], bottom=axesRange[1],
                      right=axesRange[0]+axesRange[2], 
                      top=axesRange[1]+axesRange[3], wspace=space)
    else:
        gs = GridSpec(number, 1, fig, left=axesRange[0], bottom=axesRange[1],
                      right=axesRange[0]+axesRange[2], 
                      top=axesRange[1]+axesRange[3], hspace=space)
    axes = []
    for i in range(number):
        if horizontal:
            axes.append(fig.add_subplot(gs[0, i]))
        else:
            axes.append(fig.add_subplot(gs[i, 0]))
    
    return axes


def test_hatch(x, y, p, p_limit, hatch='.', color='k', 
               ax=None, transform=ccrs.PlateCarree(), **kwargs):
    """
    画检验的hatch
    """
    if ax is None:
        ax = plt.gca()
    
    h = ax.contourf(x, y, p, levels=[0, p_limit], 
                    colors='none', hatches=[hatch], transform=transform, **kwargs)

    return h


def axisColor(ax, left=None, right=None, which='left'):

    if left is not None:
        ax.spines['left'].set_color(left)
    if right is not None:
        ax.spines['right'].set_color(right)

    if which == 'left':
        ax.tick_params(axis='y', colors=left)
        ax.yaxis.label.set_color(left)
    elif which == 'right':
        ax.tick_params(axis='y', colors=right)
        ax.yaxis.label.set_color(right)
    

def draw_box(area: list[float], ax=None, transform=ccrs.PlateCarree(), div=1, **kwargs):
    """
    画框

    Parameters
    ----------
    area: list[float]
        方框范围[左右下上]
    ax: axes | None
        axes
    transform: ccrs | None
        投影
    div: int
        方框边的分段数，经纬度弯曲的投影下分度越高越准确
    **kwargs: dict
        传给ax.plot的参数

    Return
    ------
    P: list[matplotlib.lines.Line2D]

    """
    if ax is None:
        ax = plt.gca()

    if div < 1:
        div = 1

    x = np.linspace(area[0], area[1], div+1)
    y = np.linspace(area[2], area[3], div+1)
    P = ax.plot(
        np.concatenate([x, np.ones(div+1)*area[1], x[::-1], np.ones(div+1)*area[0]]),
        np.concatenate([np.ones(div+1)*area[2], y, np.ones(div+1)*area[3], y[::-1]]),
        transform=transform, **kwargs
    )
    return P
def draw_x_marker(lon, lat, r, linewidth=1, ax=None, transform=ccrs.PlateCarree(), div=1, **kwargs):
    """
    画x标记

    Parameters
    ----------
    lon: float
        经度
    lat: float
        纬度
    r: float
        半径
    ax: axes | None
        axes
    linewidth: float
        线宽
    transform: ccrs | None
        投影
    div: int
        方框边的分段数，经纬度弯曲的投影下分度越高越准确
    **kwargs: dict
        传给ax.plot的参数

    Return
    ------
    P: list[matplotlib.lines.Line2D]

    """
    if ax is None:
        ax = plt.gca()

    if div < 1:
        div = 1

    x = np.linspace(lon-r, lon+r, div+1)
    y = np.linspace(lat-r, lat+r, div+1)
    ax.plot( x, y, transform=transform, linewidth=linewidth, **kwargs)
    ax.plot( x[::-1], y, transform=transform, linewidth=linewidth, **kwargs)


def get_marker(name: str):
    from matplotlib.path import Path
    marker = {
        "x": Path([(-1, -1), (1, 1), (1, -1), (-1, 1)] , np.array([1, 2, 1, 2], dtype=np.uint8)),
        "+" : Path([(-2.828, 0), (2.828, 0), (0, 2.828), (0, -2.828)] , np.array([1, 2, 1, 2], dtype=np.uint8)),
    }
    return marker[name]


def draw_mrg_analyse(ax: Axes, u: xr.DataArray, v: xr.DataArray, 
                     y0=6, Y=17, err_lim=0.25, minWidth=3, 
                     mrgColor='r', mrgLinewidth=0.8, vorSize=10,
                     vorMaxColor='r', vorMinColor='lime', markerlw=0.8, zorder=30):
    mrg = cal.identify_MRG(v, y0, Y, err_lim, minWidth)
    for k in mrg:
        draw_box([k[0], k[1], -17, +17], ax, color=mrgColor, linewidth=mrgLinewidth)

    # vorMax, vorMin = cal.find_local_vor(u, v, smoothRadius=10)
    # ax.scatter(vorMax[:, 0], vorMax[:, 1], s=vorSize, c=vorMaxColor, marker='^',
    #            transform=ccrs.PlateCarree(), linewidths=markerlw, zorder=zorder)
    # ax.scatter(vorMin[:, 0], vorMin[:, 1], s=vorSize, c=vorMinColor, marker='v',
    #            transform=ccrs.PlateCarree(), linewidths=markerlw, zorder=zorder)


